package publisher.confirms;

import com.rabbitmq.client.Channel;
import com.rabbitmq.client.ConfirmListener;
import com.rabbitmq.client.Connection;
import com.rabbitmq.client.ConnectionFactory;
import constant.Constants;

import java.io.IOException;
import java.util.Collections;
import java.util.SortedSet;
import java.util.TreeSet;
import java.util.concurrent.TimeoutException;

/**
 * Created with IntelliJ IDEA.
 * Description：
 *
 * @author: zxj
 * @date: 2024-09-19
 * @time: 14:08:45
 */
public class Main {

    static final int MESSAGE_COUNT = 500;

    public static void main(String[] args) throws IOException, TimeoutException, InterruptedException {
        publishMessagesIndividually();
        publishingMessagesInBatches();
        handlingPublisherConfirmsAsynchronously();
    }


    /*
     * 异步确认
     */
    private static void handlingPublisherConfirmsAsynchronously() throws IOException, TimeoutException, InterruptedException {
        try (Connection connection = getConnection()) {
            Channel channel = connection.createChannel();
            // 开启信道确认模式
            channel.confirmSelect();
            channel.queueDeclare(Constants.CONFIRM_QUEUE_NAME3, true, false, false, null);
            //有序集合,元素按照⾃然顺序进⾏排序,存储未confirm消息序号
            SortedSet<Long> confirmSet = Collections.synchronizedSortedSet(new TreeSet<>()); // 线程安全
            channel.addConfirmListener(new ConfirmListener() {
                @Override
                public void handleAck(long deliveryTag, boolean multiple) throws IOException {
                    // multiple 表示是否批量
                    // confirmSet.headSet(n)⽅法返回当前集合中⼩于n的集合
                    if (multiple) {
                        //批量确认:将集合中⼩于等于当前序号deliveryTag元素的集合清除，表⽰这批序号的消息都已经被ack了
                        confirmSet.headSet(deliveryTag + 1).clear();
                    } else {
                        //单条确认:将当前的deliveryTag从集合中移除
                        confirmSet.remove(deliveryTag);
                    }
                }

                @Override
                public void handleNack(long deliveryTag, boolean multiple) throws IOException {
                    System.err.format("deliveryTag: %d, multiple: %b%n", deliveryTag, multiple);
                    if (multiple) {
                        //批量确认:将集合中⼩于等于当前序号deliveryTag元素的集合清除，表⽰这批序号的消息都已经被ack了
                        confirmSet.headSet(deliveryTag + 1).clear();
                    } else {
                        //单条确认:将当前的deliveryTag从集合中移除
                        confirmSet.remove(deliveryTag);
                    }
                    //如果处理失败, 这⾥需要添加处理消息重发的场景. 此处代码省略
                }
            });

            // 循环发送消息

            long begin = System.currentTimeMillis();
            for (int i = 0; i < MESSAGE_COUNT; i++) {
                String msg = "消息" + i;
                // 得到下次发送信息的序号, 从 1 开始
                long nextPublishSeqNo = channel.getNextPublishSeqNo();
                // 将序号存入集合中
                confirmSet.add(nextPublishSeqNo);
                channel.basicPublish("", Constants.CONFIRM_QUEUE_NAME3, null, msg.getBytes());
            }

            // 消息发送完毕, 等待确认
            while (!confirmSet.isEmpty()) {
                Thread.sleep(10);
            }

            long end = System.currentTimeMillis();
            System.out.printf("异步确认机制: 发送 %d 条消息, 一共耗时 %d ms\n", MESSAGE_COUNT, end - begin);
        }
    }

    /*
     * 批量确认
     **/
    private static void publishingMessagesInBatches() throws IOException, TimeoutException, InterruptedException {
        try (Connection connection = getConnection()) {
            Channel channel = connection.createChannel();
            // 开启信道确认模式
            channel.confirmSelect();
            channel.queueDeclare(Constants.CONFIRM_QUEUE_NAME2, true, false, false, null);
            int havePublishCount = 0;
            int timeToConfirmCount = 100;

            long begin = System.currentTimeMillis();
            for (int i = 0; i < MESSAGE_COUNT; i++) {
                String msg = "消息" + i;
                channel.basicPublish("", Constants.CONFIRM_QUEUE_NAME2, null, msg.getBytes());
                havePublishCount++;
                if (havePublishCount == timeToConfirmCount) {
                    channel.waitForConfirms(5_000);
                    havePublishCount = 0;
                }
            }

            if (havePublishCount > 0) {
                channel.waitForConfirms(5_000);
                havePublishCount = 0;
            }
            long end = System.currentTimeMillis();
            System.out.printf("批量确认机制: 发送 %d 条消息, 一共耗时 %d ms\n", MESSAGE_COUNT, end - begin);
        }
    }

    /*
     * 单独确认
     **/
    public static void publishMessagesIndividually() throws IOException, TimeoutException, InterruptedException {
        try (Connection connection = getConnection()) {
            Channel channel = connection.createChannel();
            // 开启信道确认模式
            channel.confirmSelect();
            // 声明队列
            channel.queueDeclare(Constants.CONFIRM_QUEUE_NAME1, true, false, false, null);

            long begin = System.currentTimeMillis();
            // 循环发送消息
            for (int i = 0; i < MESSAGE_COUNT; i++) {
                // 发送消息
                String msg = "消息" + i;
                channel.basicPublish("", Constants.CONFIRM_QUEUE_NAME1, null, msg.getBytes());
                //等待确认消息.只要消息被确认，这个⽅法就会被返回
                //如果超时过期, 则抛出TimeoutException。如果任何消息被nack(丢失）,waitForConfirmsOrDie将抛出IOException。
                channel.waitForConfirms(5_000);
            }
            long end = System.currentTimeMillis();
            System.out.printf("单独确认机制: 发送 %d 条消息, 一共耗时 %d ms\n", MESSAGE_COUNT, end - begin);
        }
    }

    private static Connection getConnection() throws IOException, TimeoutException {
        // 1. 创建工厂
        ConnectionFactory factory = new ConnectionFactory();
        // 2. 设置参数
        factory.setHost(Constants.HOST);
        factory.setPort(Constants.PORT);
        factory.setVirtualHost(Constants.VIRTUAL_HOST);
        factory.setUsername(Constants.USERNAME);
        factory.setPassword(Constants.PASSWORD);
        // 3. 创建连接Connection
        return factory.newConnection();
    }
}
